Option Strict On
Option Explicit On

Imports System
Imports KadGen.MetadataWrapper
Imports System.Data
Imports System.Collections.Generic
Imports System.Data.SqlClient

Friend Class ExtendedPropertyProvider
   Private mConnection As SqlConnection
   Private mSkipExtendedProperties As Boolean
   Private mCustomAttributeNamespace As String
   Private mCustomAttributePrefix As String

   Private mDataTableCache As DataTableCache

   Friend Sub New( _
            ByVal connection As SqlConnection, _
            ByVal skipExtendedProperties As Boolean, _
            ByVal customAttributeNamespace As String, _
            ByVal customAttributePrefix As String)
      mConnection = connection
      mSkipExtendedProperties = skipExtendedProperties
      mCustomAttributeNamespace = customAttributeNamespace
      mCustomAttributePrefix = customAttributePrefix
      mDataTableCache = New DataTableCache(connection)
   End Sub

   Friend Sub AddExtendedDatabaseProperties( _
            ByVal dbsItem As WrapperBase)
      ' Get extended properties
      If Not mSkipExtendedProperties Then
         Dim extendedTable As DataTable = mDataTableCache.GetCachedTable()
         AddExtendedProperties(dbsItem, extendedTable, "")
      End If
   End Sub

   Friend Sub AddExtendedTableProperties( _
           ByVal dbsItem As WrapperBase, _
           ByVal schema As String, _
           ByVal tableName As String)
      ' Get extended properties
      If Not mSkipExtendedProperties Then
         Dim extendedTable As DataTable = mDataTableCache.GetCachedTable(schema, "TABLE")
         AddExtendedProperties(dbsItem, extendedTable, "objname = '" & tableName & "'")
      End If
   End Sub

   Friend Sub AddExtendedStoredProcProperties( _
           ByVal dbsItem As WrapperBase, _
           ByVal schema As String, _
           ByVal storedProcName As String)
      ' Get extended properties
      If Not mSkipExtendedProperties Then
         Dim extendedTable As DataTable = mDataTableCache.GetCachedTable(schema, "PROCEDURE")
         AddExtendedProperties(dbsItem, extendedTable, "objname = '" & storedProcName & "'")
      End If
   End Sub

   Friend Sub AddExtendedViewProperties( _
           ByVal dbsItem As WrapperBase, _
           ByVal schema As String, _
           ByVal viewName As String)
      ' Get extended properties
      If Not mSkipExtendedProperties Then
         Dim extendedTable As DataTable = mDataTableCache.GetCachedTable(schema, "VIEW")
         AddExtendedProperties(dbsItem, extendedTable, "objname = '" & viewName & "'")
      End If
   End Sub

   Friend Sub AddExtendedFunctionProperties( _
           ByVal dbsItem As WrapperBase, _
           ByVal schema As String, _
           ByVal functionName As String)
      ' Get extended properties
      If Not mSkipExtendedProperties Then
         Dim extendedTable As DataTable = mDataTableCache.GetCachedTable(schema, "FUNCTION")
         AddExtendedProperties(dbsItem, extendedTable, "objname = '" & functionName & "'")
      End If
   End Sub

   Friend Sub AddExtendedTableColumnProperties( _
          ByVal dbsItem As WrapperBase, _
          ByVal schema As String, _
          ByVal tableName As String, _
          ByVal columnName As String)
      ' Get extended properties
      If Not mSkipExtendedProperties Then
         Dim extendedTable As DataTable = mDataTableCache.GetCachedTable(schema, "TABLE", tableName, "COLUMN")
         AddExtendedProperties(dbsItem, extendedTable, "objname = '" & columnName & "'")
      End If
   End Sub

   Friend Sub AddExtendedViewColumnProperties( _
           ByVal dbsItem As WrapperBase, _
           ByVal schema As String, _
           ByVal viewName As String, _
           ByVal columnName As String)
      ' Get extended properties
      If Not mSkipExtendedProperties Then
         Dim extendedTable As DataTable = mDataTableCache.GetCachedTable(schema, "VIEW", viewName, "COLUMN")
         AddExtendedProperties(dbsItem, extendedTable, "objname = '" & columnName & "'")
      End If
   End Sub

   Friend Sub AddExtendedStoredProcParameterProperties( _
           ByVal dbsItem As WrapperBase, _
           ByVal schema As String, _
           ByVal storedProcName As String, _
           ByVal parameterName As String)
      ' Get extended properties
      If Not mSkipExtendedProperties Then
         Dim extendedTable As DataTable = mDataTableCache.GetCachedTable(schema, "PROCEDURE", storedProcName, "PARAMETER")
         AddExtendedProperties(dbsItem, extendedTable, "objname = '" & parameterName & "'")
      End If
   End Sub

   Friend Sub AddExtendedFunctionParameterProperties( _
           ByVal dbsItem As WrapperBase, _
           ByVal schema As String, _
           ByVal functionName As String, _
           ByVal parameterName As String)
      ' Get extended properties
      If Not mSkipExtendedProperties Then
         Dim extendedTable As DataTable = mDataTableCache.GetCachedTable(schema, "FUNCTION", functionName, "PARAMETER")
         AddExtendedProperties(dbsItem, extendedTable, "objname = '" & parameterName & "'")
      End If
   End Sub

   Private Sub AddExtendedProperties( _
               ByVal dbsItem As WrapperBase, _
               ByVal extendedTable As DataTable, _
               ByVal filter As String)
      ' Get extended properties
      If extendedTable.Rows.Count > 0 Then
         Dim extendedRows As DataRow() = extendedTable.Select(filter)
         For Each row As DataRow In extendedRows
            dbsItem.AddCustomAttribute(mCustomAttributeNamespace, _
                     mCustomAttributePrefix, row("name").ToString, _
                     row("value").ToString)
         Next
      End If
   End Sub

   Private Class DataTableCache
      Dim cache As New Dictionary(Of String, DataTable)
      Dim mConnection As SqlConnection

      Friend Sub New(ByVal connection As SqlConnection)
         mConnection = connection
      End Sub

      Public Function GetCachedTable() As DataTable
         Dim key As String = "**********"
         If cache.ContainsKey(key) Then
            Return cache(key)
         End If
         Dim dataTable As DataTable = DataTableFromSql( _
               "SELECT * FROM ::fn_listextendedproperty(null, null, null, null, null, null, null)")
         cache.Add(key, dataTable)
         Return dataTable
      End Function

      Public Function GetCachedTable( _
               ByVal schema As String, _
               ByVal type As String) _
               As DataTable
         Dim key As String = "**" & schema & "**" & type & "******"
         If cache.ContainsKey(key) Then
            Return cache(key)
         End If
         Dim dataTable As DataTable = DataTableFromSql( _
               "SELECT * FROM ::fn_listextendedproperty(null, 'User', '" & schema & _
               "', '" & type & "', null, null, null)")
         cache.Add(key, dataTable)
         Return dataTable
      End Function

      Public Function GetCachedTable( _
            ByVal schema As String, _
            ByVal type As String, _
            ByVal objectName As String, _
            ByVal level2Type As String) _
            As DataTable
         Dim key As String = "**" & schema & "**" & type & "**" & objectName & "**" & level2Type & "**"
         If cache.ContainsKey(key) Then
            Return cache(key)
         End If
         Dim dataTable As DataTable = DataTableFromSql( _
              "SELECT * FROM ::fn_listextendedproperty(null, 'User', '" & schema & _
              "', '" & type & "', '" & objectName & "', '" & level2Type & "', null)")
         cache.Add(key, dataTable)
         Return dataTable
      End Function

      Private Function DataTableFromSql(ByVal commandText As String) As DataTable
         Dim datatable As New DataTable()
         Dim dataAdapter As New SqlDataAdapter( _
                  commandText, _
                  mConnection)
         dataAdapter.Fill(datatable)
         Return datatable
      End Function



   End Class

End Class
